package com.googlecode.protoserv;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.google.gson.stream.JsonReader;
import com.googlecode.protoserv.Protocol.RequestProto;
import com.googlecode.protoserv.Protocol.ResponseProto;
import com.googlecode.protoserv.Transport.Encoding;

@SuppressWarnings("serial")
public class RpcServlet extends HttpServlet implements HttpConstants {

  public static final String PROTOSERV_RPC_MANAGER_KEY = "com.googlecode.protoserv.RpcManager";

  public static String getRpcManageKeyForPath(String path) {
    return PROTOSERV_RPC_MANAGER_KEY + ":" + path;
  }

  public RpcServlet() {
  }

  private RpcManager getRpcManager(String path) {
    return (RpcManager) getServletContext().getAttribute(
        getRpcManageKeyForPath(path));
  }

  private RequestProto readRequest(RpcManager rpcman, Encoding encoding,
      String contentType, HttpServletRequest httpRequest,
      HttpServletResponse httpResponse) throws IOException {
    InputStream input = null;
    switch (encoding.getNumber()) {
    case Encoding.IDENTITY_VALUE:
      input = httpRequest.getInputStream();
      break;
    case Encoding.GZIP_VALUE:
      input = new GZIPInputStream(httpRequest.getInputStream());
      break;
    default:
      httpResponse.sendError(HttpServletResponse.SC_UNSUPPORTED_MEDIA_TYPE);
      return null;
    }
    RequestProto request = null;
    if (MIME_PROTOBUF.equalsIgnoreCase(contentType)) {
      request = RequestProto.parseFrom(input, rpcman.getExtensionRegistry());
    } else if (MIME_JSON.equalsIgnoreCase(contentType)) {
      InputStreamReader reader = new InputStreamReader(input);
      RequestProto.Builder builder = RequestProto.newBuilder();
      JsonFormat.parse(new JsonReader(reader), rpcman.getExtensionRegistry(),
          builder);
      request = builder.build();
    } else {
      httpResponse.sendError(HttpServletResponse.SC_UNSUPPORTED_MEDIA_TYPE);
      return null;
    }
    return request;
  }

  private void writeResponse(Encoding encoding, String contentType,
      ResponseProto response, HttpServletResponse httpResponse)
      throws IOException {
    ByteArrayOutputStream buffer = new ByteArrayOutputStream(
        response.getSerializedSize());
    OutputStream output = null;
    switch (encoding.getNumber()) {
    case Encoding.IDENTITY_VALUE:
      output = buffer;
      break;
    case Encoding.GZIP_VALUE:
      output = new GZIPOutputStream(buffer);
      break;
    default:
      httpResponse.sendError(HttpServletResponse.SC_UNSUPPORTED_MEDIA_TYPE);
      return;
    }
    if (MIME_PROTOBUF.equalsIgnoreCase(contentType)) {
      response.writeTo(output);
    } else if (MIME_JSON.equalsIgnoreCase(contentType)) {
      OutputStreamWriter writer = new OutputStreamWriter(output);
      JsonFormat.stringify(response, writer);
      writer.close();
    } else {
      httpResponse.sendError(HttpServletResponse.SC_UNSUPPORTED_MEDIA_TYPE);
      return;
    }
    output.close();
    httpResponse.setStatus(HttpServletResponse.SC_OK);
    httpResponse.setHeader(USER_AGENT, PROTOSERV_USER_AGENT_PREFIX + "/j2se");
    httpResponse.setContentType(contentType);
    httpResponse.setHeader(CONTENT_ENCODING, encoding.name().toLowerCase());
    httpResponse.setContentLength(buffer.size());
    OutputStream httpOutput = httpResponse.getOutputStream();
    buffer.writeTo(httpOutput);
    httpOutput.close();
  }

  @Override
  protected void doPost(HttpServletRequest httpRequest,
      HttpServletResponse httpResponse) throws ServletException, IOException {
    try {
      RpcManager rpcman = getRpcManager(httpRequest.getPathInfo());
      if (rpcman == null) {
        httpResponse.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
        return;
      }
      String contentType = httpRequest.getContentType();
      String httpEncoding = httpRequest.getHeader(CONTENT_ENCODING);
      Encoding encoding = (httpEncoding == null) ? Encoding.IDENTITY : Encoding
          .valueOf(httpEncoding.toUpperCase());
      if (encoding == null) {
        httpResponse.sendError(HttpServletResponse.SC_UNSUPPORTED_MEDIA_TYPE);
        return;
      }
      RequestProto request;
      try {
        request = readRequest(rpcman, encoding, contentType, httpRequest,
            httpResponse);
      } catch (IOException e) {
        httpResponse.sendError(HttpServletResponse.SC_BAD_REQUEST);
        return;
      }
      try {
        ResponseProto.Builder response = ResponseProto.newBuilder();
        rpcman.processCall(request, response);
        writeResponse(encoding, contentType, response.build(), httpResponse);
      } catch (IOException e) {
        httpResponse.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
        return;
      }
    } catch (Throwable e) {
      e.printStackTrace();
      httpResponse.sendError(HttpServletResponse.SC_BAD_REQUEST);
    }
  }
}
